package com.tjmc.nsia.core.utils;

import cn.hutool.core.collection.CollectionUtil;
import com.alibaba.fastjson2.JSON;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.PropertyNamingStrategy;
import com.tjmc.nsia.core.constant.Constants;
import com.tjmc.nsia.core.dao.entity.NsiaChat;
import com.tjmc.nsia.core.domain.AjaxResult;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.service.v4.model.*;
import io.reactivex.Flowable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.StringUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * @author Luc
 * @version 1.0
 * @date 2024-06-21 11:28
 */
public class BigmodelUtils {

    private static final Logger logger = LoggerFactory.getLogger(BigmodelUtils.class);

    private static final ClientV4 client = new ClientV4.Builder(Constants.BIGMODEL_API_SECRET_KEY)
            .enableTokenCache()
            .networkConfig(300, 100, 100, 100, TimeUnit.SECONDS)
            .connectionPool(new okhttp3.ConnectionPool(8, 1, TimeUnit.SECONDS))
            .build();


    // 请自定义自己的业务id
    private static final String requestIdTemplate = "mycompany-%d";

    private static final ObjectMapper mapper = defaultObjectMapper();


    public static ObjectMapper defaultObjectMapper() {
        ObjectMapper mapper = new ObjectMapper();
        mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
        mapper.setSerializationInclusion(JsonInclude.Include.NON_NULL);
        mapper.setPropertyNamingStrategy(PropertyNamingStrategy.SNAKE_CASE);
        return mapper;
    }

    public static String syncInvoke(String question) {
        if (!StringUtils.hasText(question)) {
            return null;
        }
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), question);
        messages.add(chatMessage);
        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(com.zhipu.oapi.Constants.ModelChatGLM4)
                .stream(Boolean.FALSE)
                .invokeMethod(com.zhipu.oapi.Constants.invokeMethod)
                .messages(messages)
                .requestId(requestId)
                .build();
        ModelApiResponse invokeModelApiResp = client.invokeModelApi(chatCompletionRequest);
        try {
            logger.info("model output:" + mapper.writeValueAsString(invokeModelApiResp));
            if (!Objects.isNull(invokeModelApiResp) && invokeModelApiResp.isSuccess()) {
                ModelData modelData = invokeModelApiResp.getData();
                if (!Objects.isNull(modelData) && !CollectionUtil.isEmpty(modelData.getChoices())) {
                    return (String) modelData.getChoices().get(0).getMessage().getContent();
                }
            }
        } catch (JsonProcessingException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String sseInvoke(NsiaChat nsiaChat, String employeeCode) throws JsonProcessingException {
        if (!Objects.isNull(nsiaChat) && !StringUtils.hasText(nsiaChat.getQuestion())) {
            SseEmitterUtils.sendMessage(employeeCode, JSON.toJSONString(AjaxResult.error("查询问题为空！")));
            return null;
        }
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), nsiaChat.getQuestion());
        messages.add(chatMessage);
        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(com.zhipu.oapi.Constants.ModelChatGLM4)
                .stream(Boolean.TRUE)
                .messages(messages)
                .requestId(requestId)
                .build();
        ModelApiResponse sseModelApiResp = client.invokeModelApi(chatCompletionRequest);
        StringBuilder stringBuilder = new StringBuilder();
        if (!Objects.isNull(sseModelApiResp) && sseModelApiResp.isSuccess()) {
            AtomicBoolean isFirst = new AtomicBoolean(true);
            ChatMessageAccumulator chatMessageAccumulator = mapStreamToAccumulator(sseModelApiResp.getFlowable())
                    .doOnNext(accumulator -> {
                        {
                            String message;
                            if (isFirst.getAndSet(false)) {
                                message = "Response: ";
//                                logger.info(message);
                                stringBuilder.append(message);
                                nsiaChat.setAnswer(message);
                                if(StringUtils.hasText(employeeCode)) {
                                    SseEmitterUtils.sendMessage(employeeCode, JSON.toJSONString(AjaxResult.success(nsiaChat)));
                                }
                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getTool_calls() != null) {
                                message = mapper.writeValueAsString(accumulator.getDelta().getTool_calls());
//                                logger.info("tool_calls: " + message);
                                stringBuilder.append("tool_calls: " + message);
                                nsiaChat.setAnswer("tool_calls: " + message);
                                if(StringUtils.hasText(employeeCode)) {
                                    SseEmitterUtils.sendMessage(employeeCode, JSON.toJSONString(AjaxResult.success(nsiaChat)));
                                }
                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getContent() != null) {
                                message = accumulator.getDelta().getContent();
//                                logger.info(message);
                                stringBuilder.append(message);
                                nsiaChat.setAnswer(message);
                                if(StringUtils.hasText(employeeCode)) {
                                    SseEmitterUtils.sendMessage(employeeCode, JSON.toJSONString(AjaxResult.success(nsiaChat)));
                                }
                            }
                            if(Constants.BIGMODEL_SSE_SENDMESSAGE_INTERVAL > 0) {
                                Thread.sleep(Constants.BIGMODEL_SSE_SENDMESSAGE_INTERVAL);
                            }
                        }
                    })
                    .doOnComplete(System.out::println)
                    .lastElement()
                    .blockingGet();
//            Choice choice = new Choice(chatMessageAccumulator.getChoice().getFinishReason(), 0L, chatMessageAccumulator.getDelta());
            Choice choice = new Choice();
            choice.setFinishReason(chatMessageAccumulator.getChoice().getFinishReason());
            choice.setIndex(0L);
            choice.setDelta(chatMessageAccumulator.getDelta());
            List<Choice> choices = new ArrayList<>();
            choices.add(choice);
            ModelData data = new ModelData();
            data.setChoices(choices);
            data.setUsage(chatMessageAccumulator.getUsage());
            data.setId(chatMessageAccumulator.getId());
            data.setCreated(chatMessageAccumulator.getCreated());
            data.setRequestId(chatCompletionRequest.getRequestId());
            sseModelApiResp.setFlowable(null);
            sseModelApiResp.setData(data);
        }
        logger.info("model output:" + mapper.writeValueAsString(sseModelApiResp));
        if(StringUtils.hasText(employeeCode)) {
            SseEmitterUtils.sendMessage(employeeCode, JSON.toJSONString(sseModelApiResp));
        }
        return stringBuilder.toString();
    }

    public static String webSocketInvoke(NsiaChat nsiaChat, String userId) throws JsonProcessingException {
        WebSocketChatServer webSocketChatServer = SpringUtils.getBean(WebSocketChatServer.class);
        if(Objects.isNull(webSocketChatServer)) {
            logger.info("webSocketChatServer is null");
            return null;
        }
        if (!Objects.isNull(nsiaChat) && !StringUtils.hasText(nsiaChat.getQuestion())) {
            webSocketChatServer.sendOneMessage(userId, JSON.toJSONString(AjaxResult.error("查询问题为空！")));
//            webSocketChatServer.sendMessageAsynchronously(userId, JSON.toJSONString(AjaxResult.error("查询问题为空！")));
            return null;
        }
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), nsiaChat.getQuestion());
        messages.add(chatMessage);
        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(com.zhipu.oapi.Constants.ModelChatGLM4)
                .stream(Boolean.TRUE)
                .messages(messages)
                .requestId(requestId)
                .build();
        ModelApiResponse sseModelApiResp = client.invokeModelApi(chatCompletionRequest);
        StringBuilder stringBuilder = new StringBuilder();
        if (!Objects.isNull(sseModelApiResp) && sseModelApiResp.isSuccess()) {
            AtomicBoolean isFirst = new AtomicBoolean(true);
            ChatMessageAccumulator chatMessageAccumulator = mapStreamToAccumulator(sseModelApiResp.getFlowable())
                    .doOnNext(accumulator -> {
                        {
                            String message;
                            if (isFirst.getAndSet(false)) {
                                message = "Response: ";
                                stringBuilder.append(message);
                                nsiaChat.setAnswer(message);
                                if(StringUtils.hasText(userId)) {
                                    webSocketChatServer.sendOneMessage(userId, JSON.toJSONString(AjaxResult.success(nsiaChat)));
//                                    webSocketChatServer.sendMessageAsynchronously(userId, JSON.toJSONString(AjaxResult.success(nsiaChat)));
                                }
                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getTool_calls() != null) {
                                message = mapper.writeValueAsString(accumulator.getDelta().getTool_calls());
                                stringBuilder.append("tool_calls: " + message);
                                nsiaChat.setAnswer("tool_calls: " + message);
                                if(StringUtils.hasText(userId)) {
                                    webSocketChatServer.sendOneMessage(userId, JSON.toJSONString(AjaxResult.success(nsiaChat)));
//                                    webSocketChatServer.sendMessageAsynchronously(userId, JSON.toJSONString(AjaxResult.success(nsiaChat)));
                                }
                            }
                            if (accumulator.getDelta() != null && accumulator.getDelta().getContent() != null) {
                                message = accumulator.getDelta().getContent();
                                stringBuilder.append(message);
                                nsiaChat.setAnswer(message);
                                if(StringUtils.hasText(userId)) {
                                    webSocketChatServer.sendOneMessage(userId, JSON.toJSONString(AjaxResult.success(nsiaChat)));
//                                    webSocketChatServer.sendMessageAsynchronously(userId, JSON.toJSONString(AjaxResult.success(nsiaChat)));
                                }
                            }
                            if(Constants.BIGMODEL_SSE_SENDMESSAGE_INTERVAL > 0) {
                                Thread.sleep(Constants.BIGMODEL_SSE_SENDMESSAGE_INTERVAL);
                            }
                        }
                    })
                    .doOnComplete(System.out::println)
                    .lastElement()
                    .blockingGet();
            Choice choice = new Choice();
            choice.setFinishReason(chatMessageAccumulator.getChoice().getFinishReason());
            choice.setIndex(0L);
            choice.setDelta(chatMessageAccumulator.getDelta());
            List<Choice> choices = new ArrayList<>();
            choices.add(choice);
            ModelData data = new ModelData();
            data.setChoices(choices);
            data.setUsage(chatMessageAccumulator.getUsage());
            data.setId(chatMessageAccumulator.getId());
            data.setCreated(chatMessageAccumulator.getCreated());
            data.setRequestId(chatCompletionRequest.getRequestId());
            sseModelApiResp.setFlowable(null);
            sseModelApiResp.setData(data);
        }
        logger.info("model output:" + mapper.writeValueAsString(sseModelApiResp));
        if(StringUtils.hasText(userId)) {
            webSocketChatServer.sendOneMessage(userId, JSON.toJSONString(AjaxResult.success(sseModelApiResp.getData())));
//            webSocketChatServer.sendMessageAsynchronously(userId, JSON.toJSONString(AjaxResult.success(sseModelApiResp.getData())));
        }
        return stringBuilder.toString();
    }

    private static Flowable<ChatMessageAccumulator> mapStreamToAccumulator(Flowable<ModelData> flowable) {
        return flowable.map(chunk -> new ChatMessageAccumulator(chunk.getChoices().get(0).getDelta(), null, chunk.getChoices().get(0), chunk.getUsage(), chunk.getCreated(), chunk.getId()));
    }
}
